import yaml
import resource
import os
import sys
import shlex, subprocess
from threading import Thread

config_full_file_path = 'sandbox.yml'

class TimedProgram(Thread):
    """based on http://code.activestate.com/recipes/483752-timelimit-tell-a-function-to-time-out-after-a-time/"""
    def __init__(self,prog,config):
        Thread.__init__(self)
        self.exit_code = None
        self.prog = prog
        self.config = config
    
    def run(self):
        stdin, stdout, stderr = open_streams(self.config['stream'])
        args = shlex.split(self.prog)
        print args
        self.exit_code = subprocess.call(args=args,stdin=stdin, stdout=stdout, stderr=stderr)
        close_streams(stdin,stdout,stderr)

def set_resource_limits(limit):
    """Setting the resouce limits for the 3rd party processes that will run"""
    #sets time limit
    #TODO: fix problem of no placing time limit on sleeping threads
    resource.setrlimit(resource.RLIMIT_CPU,(limit['tsec'],limit['tsec']))
    #sets heap limit
    max_mem = limit['memM'] * 1024 * 1024
    #both resource limits needed to set memory limit
    resource.setrlimit(resource.RLIMIT_DATA,(max_mem,max_mem))
    resource.setrlimit(resource.RLIMIT_AS,(max_mem,max_mem))
    #setting the max number of threads/forks
    #TODO: test with value of 1
    min_process_count = min(limit['thread'],limit['fork'])
    resource.setrlimit(resource.RLIMIT_NPROC,(min_process_count,min_process_count))

def get_file_handler(handler,mode='r'):
    if handler == 'stdin': 
        sys.stdin.flush()
        return sys.stdin
    elif handler == 'stdout': 
        sys.stdout.flush()
        return sys.stdout
    elif handler == 'stderr': 
        sys.stderr.flush()
        return sys.stderr
    else: return open(handler,mode)
    
def open_streams(streams):
    """returns file objects of the redirection streams based on the config file"""
    return (get_file_handler(streams['stdin'],'r'),get_file_handler(streams['stdout'],'w'),get_file_handler(streams['stderr'],'w'))

def close_streams(stdin,stdout,stderr):
    if stdin==sys.stdin: stdin.flush()
    else: stdin.close()
    if stdout==sys.stdout: stdout.flush()
    else: stdout.close()
    if stderr==sys.stderr: stderr.flush()
    else: stderr.close()
    
if __name__=='__main__':
    f = open(config_full_file_path)
    config = yaml.load(f.read())
    f.close()
    #print `config`
    #creating chroot
    #TODO: fix hack for breaking out of chroot jail found in http://www.bpfh.net/simes/computing/chroot-break.html
    os.chdir(config['root'])
    #os.chroot(config['root'])
    #setting resource limits
    set_resource_limits(config['limit'])
    #running processes
    for prog in config['exec']:
        process = TimedProgram(prog,config)
        process.start()
        process.join(config['limit']['tsec'])
        if process.isAlive():
            #print "program still alive"
            print -9
            process._Thread__stop()
        else:
            print process.exit_code
